import numpy as np
from collections import Counter

class KNN:
    def __init__(self, k=3):
        """初始化KNN模型，指定近邻数量k"""
        self.k = k
        self.X_train = None  # 训练特征
        self.y_train = None  # 训练标签

    def fit(self, X, y):
        """训练模型（KNN是惰性学习，仅存储训练数据）"""
        self.X_train = X
        self.y_train = y

    def _euclidean_distance(self, x1, x2):
        """计算两个样本之间的欧氏距离"""
        return np.sqrt(np.sum((x1 - x2) **2))

    def predict(self, X):
        """预测新样本的类别"""
        predictions = [self._predict_single(x) for x in X]
        return np.array(predictions)

    def _predict_single(self, x):
        """预测单个样本的类别"""
        # 1. 计算与所有训练样本的距离
        distances = [self._euclidean_distance(x, x_train) for x_train in self.X_train]
        
        # 2. 按距离排序，取前k个样本的索引
        k_indices = np.argsort(distances)[:self.k]
        
        # 3. 取前k个样本的标签
        k_nearest_labels = [self.y_train[i] for i in k_indices]
        
        # 4. 多数投票决定预测结果
        most_common = Counter(k_nearest_labels).most_common(1)
        return most_common[0][0]


# 测试代码
if __name__ == "__main__":
    # 生成示例数据（分类问题：两个类别）
    X_train = np.array([
        [1, 2], [2, 3], [3, 4], [6, 7], [7, 8], [8, 9]  # 特征
    ])
    y_train = np.array([0, 0, 0, 1, 1, 1])  # 标签（0或1）

    # 初始化并训练模型
    knn = KNN(k=3)
    knn.fit(X_train, y_train)

    # 预测新样本
    X_test = np.array([[4, 5], [5, 6]])  # 待预测样本
    predictions = knn.predict(X_test)

    print("预测结果：", predictions)  # 输出：[0 1]（根据距离判断）